import cv2 as cv
import sys
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import numpy as np
import time
import datetime


def my_show_img(img, title="no title", trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)
    plt.axis('off')
    plt.title(title)


# img_dir = '../../../../../large_data/CV2/lesson/Day07'
# img_name = 'football_ball.jpg'
# temp_path = os.path.join(img_dir, img_name)

# img_dir = '../../../../../large_data/pic/'
# img_name = 'DSC05039.JPG'
# img_name = 'DSC05022_1.JPG'
# img_name = 'dog.jpg'
# img_name = 'dog_bird.jpg'
# img_name = 'football.jpg'
# img_name = 'messi5.jpg'
img_dir = '../../../../../large_data/pic/watershed/'
img_name = 'water_coins.jpg'
img_path = os.path.join(img_dir, img_name)

spr = 3
spc = 5
spn = 0
plt.figure(figsize=[13, 7])

sep('load')
img = cv.imread(img_path, cv.IMREAD_COLOR)
print('original shape', img.shape)
H, W, CH = img.shape
H2 = H // 2
W2 = W // 2
my_show_img(img, 'original', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

sep('gray')
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
my_show_img(gray, 'gray', cmap='gray')

sep('bin')
ret, thresh = cv.threshold(gray, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)
my_show_img(thresh, 'otsu', cmap='gray')

sep('noise removal')
kernel = np.ones((3, 3), np.uint8)
opening = cv.morphologyEx(thresh, cv.MORPH_OPEN, kernel, iterations=2)
my_show_img(opening, 'bin => opening', cmap='gray')

sep('sure bg')
sure_bg = cv.dilate(opening, kernel, iterations=9)
print_numpy_ndarray_info(sure_bg, 'sure_bg')
my_show_img(sure_bg, 'opening => sure_bg', cmap='gray')

sep('sure fg')
dist = cv.distanceTransform(opening, cv.DIST_L2, 5)
print_numpy_ndarray_info(dist, 'dist')
my_show_img(dist, 'dist', cmap='gray')
dist_max = dist.max()
ret, sure_fg = cv.threshold(dist, 0.7*dist_max, dist_max, cv.THRESH_BINARY)
my_show_img(sure_fg, 'sure_fg', cmap='gray')

sep('unknown region')
cv.normalize(sure_fg, sure_fg, 0, 255, cv.NORM_MINMAX)
sure_fg = np.uint8(sure_fg)
unknown = cv.subtract(sure_bg, sure_fg)
my_show_img(unknown, 'unknown', cmap='gray')

sep('labelling')
markers = cv.bitwise_not(unknown)
print_numpy_ndarray_info(markers, 'markers')
my_show_img(markers, 'markers', cmap='gray')
ret, markers = cv.connectedComponents(markers)
print_numpy_ndarray_info(markers, 'markers')
my_show_img(markers, 'markers')

sep('watershed')
markers = cv.watershed(img, markers)
my_show_img(img, 'after', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
# my_show_img(markers, 'markers', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))  # Unsupported depth of input image

colorTab = [(255, 255, 255)]
for i in range(len(np.unique(markers))):
    b = np.random.randint(0, 256)
    g = np.random.randint(0, 256)
    r = np.random.randint(0, 256)
    colorTab.append((b, g, r))
colorTab = np.uint8(colorTab)
# color table end
markers[markers == -1] = 0
markers_show = colorTab[markers]
my_show_img(markers_show, 'markers_show', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

img[markers == 0] = [0, 255, 0]
my_show_img(img, 'after ticked', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))
